package com.yjz.media.study.util.camera;

import android.graphics.ImageFormat;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Bundle;
import android.util.Size;

import androidx.annotation.NonNull;
import androidx.camera.core.ImageAnalysis;
import androidx.camera.core.ImageProxy;

import com.blankj.utilcode.util.FileIOUtils;
import com.blankj.utilcode.util.LogUtils;
import com.blankj.utilcode.util.ToastUtils;
import com.yjz.media.study.MediaApplication;
import com.yjz.media.study.cons.RTMPConfigure;
import com.yjz.media.study.jni.X264NativeHelper;
import com.yjz.media.study.rtmp.RTMPPacket;
import com.yjz.media.study.rtmp.RtmpLivePush;
import com.yjz.media.study.util.ByteUtils;
import com.yjz.media.study.util.YuvUtils;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 作用 : cameraX 数据分析
 * 作者 : yjz
 * 时间 : 2023/11/8 11:16
 */
public class CameraImageByteAnalyzer implements ImageAnalysis.Analyzer {

    // 线程锁
    private final ReentrantLock lock = new ReentrantLock();
    // 音解码器
    private MediaCodec mediaCodec = null;
    // yuv420的yuv数据
    private byte[] y;
    private byte[] u;
    private byte[] v;
    // yuvToNv21的数据
    private byte[] nv21;
    // 旋转的数据
    private byte[] nv21_rotated;
    // 将数据回调的接口
    private IPushDataListener pushDataListener;

    // x264库初始化的状态
    private boolean x264Init;
    // 是否支持硬解码
    private boolean dspSupport;

    // x264重试次数
    private int retryCount;
    // 用于记录实际，超过2s需要编码一个I帧
    private long timeStamp;

    public CameraImageByteAnalyzer() {
        dspSupport = RTMPConfigure.IS_USE_DSP;
    }

    /**
     * 停止推流，释放资源
     */
    public void stopPush() {
        if (mediaCodec != null) {
            timeStamp = 0;
            mediaCodec.stop();
            mediaCodec.release();
            mediaCodec = null;
        }
    }

    @Override
    public void analyze(@NonNull ImageProxy image) {
        if (RtmpLivePush.getInstance().isRunning()) {
            try { // 防止释放mediaCodec时，mediaCodec还在使用造成异常
                // Y:U:V 4:2:2
                if (image.getFormat() == ImageFormat.YUV_420_888) {
                    ImageProxy.PlaneProxy[] planes = image.getPlanes();
                    // 加锁确保y、u、v来源于同一个Image
                    lock.lock();
                    if (y == null || u == null || v == null) {
                        // 重复使用同一批byte数组，减少gc频率
                        y = new byte[planes[0].getBuffer().limit() - planes[0].getBuffer().position()];
                        u = new byte[planes[1].getBuffer().limit() - planes[1].getBuffer().position()];
                        v = new byte[planes[2].getBuffer().limit() - planes[2].getBuffer().position()];
                    }

                    if (image.getPlanes()[0].getBuffer().remaining() == y.length) {
                        planes[0].getBuffer().get(y);
                        planes[1].getBuffer().get(u);
                        planes[2].getBuffer().get(v);
                        Size size = new Size(image.getWidth(), image.getHeight());
                        int width = size.getHeight();
                        int height = image.getWidth();
                        if (nv21 == null) {
                            nv21 = new byte[height * width * 3 / 2];
                            nv21_rotated = new byte[height * width * 3 / 2];
                        }

                        if (RTMPConfigure.IS_USE_DSP) { // 是否使用硬解码
                            if (mediaCodec == null && dspSupport) {
                                initCodec(size);
                            }
                        }

                        // 将YUV_420_888数据转换成nv21
                        YuvUtils.yuvToNv21(y, u, v, nv21, height, width);
                        // 旋转nv21数据
                        YuvUtils.nv21_rotate_to_90(nv21, nv21_rotated, height, width);

                        if (!dspSupport) { // 硬编码不支持
                            if (pushDataListener != null) {
                                // 超过最大的重试次数
                                if (retryCount > RTMPConfigure.MAX_RETRY_COUNT) {
                                    ToastUtils.showShort("X264初始化失败！");
                                    lock.unlock();
                                    return;
                                }
                                if (!x264Init) {
                                    LogUtils.eTag("rtmp", "x264 Init");
                                    x264Init = X264NativeHelper.getInstance().init(size.getHeight(), size.getWidth(), 10, 640_000);
                                    // 初始化失败，进行重试
                                    if (!x264Init) {
                                        retryCount++;
                                        return;
                                    }
                                    LogUtils.eTag("rtmp", "x264 Init Success!");
                                }
                                // x264 初始化成功时才push数据
                                RTMPPacket rtmpPacket = new RTMPPacket(nv21_rotated, RTMPPacket.VIDE_TYPE);
                                pushDataListener.videoDataPush(rtmpPacket, false);
                            }
                        } else { // 使用硬编码
                            if (System.currentTimeMillis() - timeStamp > 2000) { // 2s触发一次I帧
                                Bundle bundle = new Bundle();
                                bundle.putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 1);
                                mediaCodec.setParameters(bundle);
                                timeStamp = System.currentTimeMillis();
                            }

                            byte[] nv12 = YuvUtils.nv21toNV12(nv21_rotated);
                            // 将要编码的数据压入
                            int inIndex = mediaCodec.dequeueInputBuffer(100_000);
                            if (inIndex >= 0) {
                                ByteBuffer byteBuffer = mediaCodec.getInputBuffer(inIndex);
                                byteBuffer.clear();
                                byteBuffer.put(nv12, 0, nv12.length);
                                mediaCodec.queueInputBuffer(inIndex, 0, nv12.length, 0, 0);
                            }
                            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
                            int outIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, 100_000);
                            if (outIndex >= 0) {
                                ByteBuffer outputBuffer = mediaCodec.getOutputBuffer(outIndex);
                                byte[] bytes = new byte[outputBuffer.remaining()];
                                outputBuffer.get(bytes);
                                if (pushDataListener != null) {
                                    RTMPPacket rtmpPacket = new RTMPPacket(bytes, RTMPPacket.VIDE_TYPE);
                                    pushDataListener.videoDataPush(rtmpPacket, true);
                                }
                                String hexStr = ByteUtils.hexBytesToString(bytes);
                                LogUtils.eTag("video", "hexStr = ", hexStr);
                                // h264文件
                                FileIOUtils.writeFileFromBytesByStream(MediaApplication.getContext().getExternalCacheDir() + "/camerax_h264.h264", bytes, true);
                                // h264数据
                                FileIOUtils.writeFileFromString(MediaApplication.getContext().getExternalCacheDir() + "/camerax_h264.txt", hexStr, true);
                                mediaCodec.releaseOutputBuffer(outIndex, false);
                            }
                        }
                    }
                    lock.unlock();
                }
                // 需要close，不然只会回调一次调用
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        image.close();
    }

    /**
     * 初始化硬解码器
     */
    private void initCodec(Size size) {
        try {
            mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
            MediaFormat mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, size.getHeight(), size.getWidth());
            //设置帧率
            mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible);
            mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 10);
            mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, 640_000);
            mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);//2s一个I帧
            mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
            mediaCodec.start();
        } catch (IOException e) { // mediaCodec出现异常时，初始化x264
            dspSupport = false;
            x264Init = X264NativeHelper.getInstance().init(size.getHeight(), size.getWidth(), 10, 640_000);
            LogUtils.eTag("rtmp", "X264NativeHelper init status " + x264Init);
        }
    }

    /**
     * CameraX数据的接口回调
     */
    public interface IPushDataListener {
        void videoDataPush(RTMPPacket rtmpPacket, boolean useDsp);
    }

    public void setPushDataListener(IPushDataListener pushDataListener) {
        this.pushDataListener = pushDataListener;
    }
}
